import networkx as nx
import numpy as np
import sys
import os
import pylab as plt
import math
import scipy.optimize as opt
from params import *

FILETYPE = 'pdf'
datasets = ['bkite','fsquare','gowalla']
lbls = ['Brightkite', 'Foursquare', 'Gowalla']
node_dist = [5651.0,8494.0,5663.0]
link_dist = [2041.0,1442.0,1792.0]
degs = [7.88,22.07,9.48,7.0]
plots = []
for dataset in datasets:
    #prob_file = os.path.join(WORKDIR, 'gsn','results',dataset, '%s_friend_prob_geo.txt'%dataset)
    #prob_file = os.path.join(WORKDIR, 'gsn','results',dataset, '%s_friend_prob_geo_10.txt'%dataset)
    prob_file = os.path.join(WORKDIR, 'gsn','results',dataset, '%s_friend_prob_geo_log.txt'%dataset)
    print dataset
    print prob_file

    values = []
    probs = []
    for line in open(prob_file):
        v,p = map(float,line.split(';'))
        values.append(v)
        probs.append(p)

    def fit_values_opt(values,probs):
        # define our (line) fitting function
        values = np.array(values)
        probs = np.array(probs)
        fitfunc = lambda p, x: p[0]*(x**(p[1])) + p[2]
        fitfunc = lambda p, x: p[0]*(x**(p[1]))
        errfunc = lambda p, x, y: (y - fitfunc(p, x))
        pinit = (1e-2,-1.0,1e-5)
        out = opt.leastsq(errfunc, pinit,args=(values, probs))

        pfinal = out[0]
        a,beta,eps = pfinal
        print a, beta, eps
        fitted = fitfunc(pfinal,values)
        return pfinal, fitted

    def fit_values(values,probs):
        values = np.log(np.array(values))
        probs = np.log(np.array(probs))
        print 'polyfit'
        p = plt.polyfit(values,probs,1)
        print 'polyval'
        fitted = plt.polyval(p,values)
        fitted = np.exp(fitted)
        a,b = p
        print a,b
        return (b,a), fitted

    MIN = 1
    MAX = 10000
    #MIN = 100
    #MAX = 1000
    values, probs
    v,p = [],[]
    for a,b in zip(values,probs):
        if MIN <= a <= MAX:
            v.append(a)
            p.append(b)
    print len(v), len(p)
    print 'fitting'
    parameters, fitted = fit_values(v,p)
    #parameters, fitted = fit_values_opt(v,p)
    a,beta = parameters[:2]
    first1 = probs[1]
    first2 = first1#1e-2
    fit1 = [v[1]*first1/x for x in v]
    fit2 = [(v[1]**0.5)*first2/(x**0.5) for x in v]
    #plots.append((values,probs,fitted,beta))
    plots.append((v,p,fitted,beta))
    continue

    print 'plotting'
    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    plt.loglog(v,p,'co')
    ax1 = plt.loglog(v,fit1,'k-')
    ax2 = plt.loglog(v,fit2,'k--')
    plt.loglog(v,fitted,'k--',linewidth=1)
    #plt.figtext(0.2,0.2,r'$P \sim d^{%.2f}$'%(beta),fontsize=24)
    #plt.figtext(0.3,0.4,r'$\alpha = %.2f$'%(beta),fontsize=16,
    #        bbox=dict(facecolor='w',fill=True,edgecolor='k'))
    plt.figtext(0.25,0.3,r'$\alpha = 1.0$',fontsize=14,
            bbox=dict(facecolor='w',fill=True,edgecolor='k'))
    plt.figtext(0.6,0.8,r'$\alpha = 0.5$',fontsize=14,
            bbox=dict(facecolor='w',fill=True,edgecolor='k'))
    #plt.legend([ax1,ax2],[r'$\alpha = 1$',r'$\alpha=0.5'],
    #        loc='lower left')
    #plt.xlim((1,20000))
    plt.xlim((1, 10))
    plt.ylim((1e-6,1e-2))
    plt.grid(True)
    plt.xlabel('Distance [km]')
    plt.ylabel('Probablity of friendship')
    plt.savefig('%s_friendprob_geo.%s'%(dataset,FILETYPE))
    plt.close()

plt.figure()
plt.clf()
plt.axes(FIG_AXES2)
FILETYPE = 'pdf'
legs = []
for d,m in zip(plots,markers):
    v,p,f,beta = d
    ax = plt.loglog(v,p,'k%s'%m, mfc='None')
    legs.append(ax[0])

plt.loglog(v,fitted,'k--',linewidth=1)
#plt.figtext(0.2,0.2,r'$P \sim d^{%.2f}$'%(beta),fontsize=24)
#plt.figtext(0.3,0.4,r'$\alpha = %.2f$'%(beta),fontsize=16,
#        bbox=dict(facecolor='w',fill=True,edgecolor='k'))
#fit1 = [1e-2/x for x in v]
#fit2 = [1e-2/(x**0.5) for x in v]
##plt.figtext(0.55,0.25,r'$\alpha = 1.0$',fontsize=12,
##        bbox=dict(facecolor='w',fill=True,edgecolor='k'))
##plt.figtext(0.7,0.8,r'$\alpha = 0.5$',fontsize=12,
##        bbox=dict(facecolor='w',fill=True,edgecolor='k'))
#ax1 = plt.loglog(v,fit1,'k-')
#ax2 = plt.loglog(v,fit2,'k-')
plt.legend(legs,lbls,
        loc='lower left', numpoints=1)
plt.xlim((MIN, MAX))
plt.ylim((1e-6,1e-2))
plt.grid(True)
plt.xlabel('Distance [km]')
plt.ylabel('Probability of friendship')
plt.savefig('three_friendprob_geo.%s'%(FILETYPE))
plt.close()
